import bpy
import sys
import os
import json
import argparse
import mathutils
import numpy as np

def import_3d_model(model_path):
    """根据文件扩展名导入不同格式的3D模型"""
    file_ext = os.path.splitext(model_path)[1].lower()
    
    print(f"正在导入 {file_ext} 格式的模型: {model_path}")
    
    try:
        if file_ext == '.obj':
            bpy.ops.wm.obj_import(filepath=model_path)
            print("OBJ 文件导入成功")
        elif file_ext == '.fbx':
            bpy.ops.import_scene.fbx(filepath=model_path)
            print("FBX 文件导入成功")
        elif file_ext in ['.glb', '.gltf']:
            bpy.ops.import_scene.gltf(filepath=model_path)
            print(f"{file_ext.upper()} 文件导入成功")
        else:
            raise ValueError(f"不支持的文件格式: {file_ext}")
            
    except Exception as e:
        print(f"{file_ext.upper()} 导入失败: {e}")
        raise

def calculate_best_projection_direction(obj, face_indices):
    """计算选中面的最佳投影方向"""
    # 计算面的平均法向
    avg_normal = mathutils.Vector((0, 0, 0))
    for idx in face_indices:
        avg_normal += obj.data.polygons[idx].normal
    avg_normal.normalize()
    
    # 找出与平均法向最接近的坐标轴
    axes = [
        mathutils.Vector((1, 0, 0)), 
        mathutils.Vector((0, 1, 0)), 
        mathutils.Vector((0, 0, 1))
    ]
    
    best_axis = None
    best_dot = -1
    for axis in axes:
        dot = abs(avg_normal.dot(axis))
        if dot > best_dot:
            best_dot = dot
            best_axis = axis
    
    print(f"最佳投影方向: {best_axis}")
    return best_axis

def project_uvs(obj, face_indices, projection_axis, uv_layer_name):
    """手动计算UV投影"""
    # 获取UV层
    uv_layer = obj.data.uv_layers[uv_layer_name]
    
    # 收集所有选中面的顶点坐标
    all_coords = []
    for face_idx in face_indices:
        face = obj.data.polygons[face_idx]
        for vert_idx in face.vertices:
            vert = obj.data.vertices[vert_idx]
            all_coords.append(vert.co)
    
    # 计算投影方向
    if projection_axis.x > 0.9:  # X轴
        u_axis = 1  # Y
        v_axis = 2  # Z
    elif projection_axis.y > 0.9:  # Y轴
        u_axis = 0  # X
        v_axis = 2  # Z
    else:  # Z轴
        u_axis = 0  # X
        v_axis = 1  # Y
    
    # 计算UV范围
    min_u = min(coord[u_axis] for coord in all_coords)
    max_u = max(coord[u_axis] for coord in all_coords)
    min_v = min(coord[v_axis] for coord in all_coords)
    max_v = max(coord[v_axis] for coord in all_coords)
    
    range_u = max_u - min_u
    range_v = max_v - min_v
    
    if range_u == 0:
        range_u = 1
    if range_v == 0:
        range_v = 1
    
    # 应用UV投影
    for face_idx in face_indices:
        face = obj.data.polygons[face_idx]
        for loop_idx in face.loop_indices:
            loop = obj.data.loops[loop_idx]
            vert = obj.data.vertices[loop.vertex_index]
            coord = vert.co
            
            # 计算UV坐标
            u = (coord[u_axis] - min_u) / range_u
            v = (coord[v_axis] - min_v) / range_v
            
            # 设置UV坐标
            uv_layer.data[loop_idx].uv = (v, u)
    
    print("手动UV投影完成")

def setup_texture_painting(
    obj_path,              # 模型文件路径
    output_path,           # 输出文件路径
    texture_path=None,
    base_color=(1.0, 1.0, 1.0, 1.0),
    metallic=0.0,  # 金属度
    roughness=0.5, # 粗糙度
    material_name=None     # 材质名称（用于部分区域选择）
):
    # 清除当前场景中的所有对象
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    
    # 检查文件是否存在
    if not os.path.exists(obj_path):
        print(f"错误: 模型文件不存在 - {obj_path}")
        return
    
    # 导入3D模型 - 支持多种格式
    try:
        import_3d_model(obj_path)
    except Exception as e:
        print(f"模型导入失败: {e}")
        return
    
    # 获取导入的对象
    obj = bpy.context.selected_objects[0] if bpy.context.selected_objects else None
    for o in bpy.context.selected_objects:
        if o.type == 'MESH':
            obj = o
            break   
        
    if not obj:
        print("错误: 没有找到有效的对象")
        return
    else:
        print(f"已找到对象: {obj.name} (类型: {obj.type})")
    
    # 选择对象
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    
    # 检查对象是否为网格
    if obj.type != 'MESH':
        print(f"错误: 对象类型不是网格 (类型: {obj.type})")
        return
    
    # 检查纹理路径
    if texture_path:
        texture_path = os.path.normpath(texture_path)
        print(f"检查纹理路径: {texture_path}")
        
        if not os.path.exists(texture_path):
            print(f"错误: 纹理文件不存在 - {texture_path}")
            texture_path = None
    
    # 如果没有有效路径，创建临时纹理
    if not texture_path:
        print("创建临时纹理...")
        temp_dir = os.path.dirname(bpy.data.filepath) if bpy.data.filepath else os.getcwd()
        texture_path = os.path.join(temp_dir, "temp_texture.png")
        
        image = bpy.data.images.new("TempTexture", width=1024, height=1024)
        pixels = [1.0, 0.5, 0.0, 1.0] * (1024 * 1024)  # 橙色
        image.pixels = pixels
        
        try:
            image.save_render(filepath=texture_path)
            print(f"临时纹理已保存到: {texture_path}")
        except Exception as e:
            print(f"无法保存临时纹理: {e}")
            texture_path = None
    
    # 部分区域选择和UV处理 ====================================================
    print("="*50)
    print("部分区域UV处理")
    
    # 查找目标材质的索引
    target_material_index = -1
    if material_name:
        print(f"查找目标材质: {material_name}")
        for i, slot in enumerate(obj.material_slots):
            if slot.material and slot.material.name == material_name:
                target_material_index = i
                print(f"找到目标材质，索引: {target_material_index}")
                break
        
        if target_material_index == -1:
            print(f"警告: 材质 '{material_name}' 不存在！")
            return
    else:
        print("警告: 未指定材质名称！")
        return
    
    # 进入编辑模式并选择面
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_mode(type='FACE')
    bpy.ops.mesh.select_all(action='DESELECT')
    
    # 选择指定材质的面
    selected_faces = []
    for poly in obj.data.polygons:
        if poly.material_index == target_material_index:
            poly.select = True
            selected_faces.append(poly.index)
    
    print(f"选中了 {len(selected_faces)} 个面")
    
    if len(selected_faces) == 0:
        print("错误: 没有找到使用指定材质的面！")
        bpy.ops.object.mode_set(mode='OBJECT')
        return
    
    # 返回对象模式
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # 创建新的UV层（只用于选中区域）
    uv_layer_name = "Texture_UV"
    if uv_layer_name in obj.data.uv_layers:
        obj.data.uv_layers.remove(obj.data.uv_layers[uv_layer_name])
    
    new_uv_layer = obj.data.uv_layers.new(name=uv_layer_name)
    obj.data.uv_layers.active = new_uv_layer
    print(f"创建新UV层: {uv_layer_name}")
    
    # 计算最佳投影方向
    projection_axis = calculate_best_projection_direction(obj, selected_faces)
    
    # 手动应用UV投影
    project_uvs(obj, selected_faces, projection_axis, uv_layer_name)
    
    print("="*50)
    
    # 创建新的纹理材质（只用于选中区域）===================================
    print("创建新的纹理材质")
    
    # 创建新材质
    texture_material = bpy.data.materials.new(name=f"{material_name}_Textured")
    texture_material.use_nodes = True
    
    # 获取材质节点
    nodes = texture_material.node_tree.nodes
    links = texture_material.node_tree.links
    
    # 清除默认节点（保留 Material Output）
    output_node = None
    for node in list(nodes):
        if node.name == "Material Output":
            output_node = node
        else:
            nodes.remove(node)
    
    if not output_node:
        output_node = nodes.new('ShaderNodeOutputMaterial')
        output_node.location = (300, 0)
    
    # 创建基础节点
    bsdf_node = nodes.new('ShaderNodeBsdfPrincipled')
    bsdf_node.location = (0, 0)
    
    # 设置基础色
    clamped_color = tuple(min(max(c, 0.0), 1.0) for c in base_color)
    bsdf_node.inputs['Base Color'].default_value = clamped_color
    
    # 设置金属度和粗糙度
    bsdf_node.inputs['Metallic'].default_value = metallic
    bsdf_node.inputs['Roughness'].default_value = roughness
    
    # 创建纹理节点
    texture_node = nodes.new('ShaderNodeTexImage')
    texture_node.location = (-300, 0)
    
    # 加载纹理
    if texture_path:
        try:
            texture_node.image = bpy.data.images.load(texture_path)
            print(f"成功加载纹理: {texture_path}")
        except Exception as e:
            print(f"加载纹理失败: {e}")
            texture_node.image = bpy.data.images.new("DefaultTexture", width=1024, height=1024)
    else:
        texture_node.image = bpy.data.images.new("FallbackTexture", width=1024, height=1024)
    
    # 创建UV贴图节点，指定使用新的UV层
    uv_map_node = nodes.new('ShaderNodeUVMap')
    uv_map_node.location = (-500, 0)
    uv_map_node.uv_map = uv_layer_name
    
    # 连接节点
    try:
        links.new(uv_map_node.outputs['UV'], texture_node.inputs['Vector'])
        links.new(texture_node.outputs['Color'], bsdf_node.inputs['Base Color'])
        links.new(bsdf_node.outputs['BSDF'], output_node.inputs['Surface'])
        print("节点连接成功")
    except Exception as e:
        print(f"节点连接失败: {e}")
        return
    
    # 将新材质应用到对象上（作为新的材质槽）
    obj.data.materials.append(texture_material)
    new_material_index = len(obj.data.materials) - 1
    print(f"新材质已添加，索引: {new_material_index}")
    
    # 将选中的面分配给新材质
    for poly in obj.data.polygons:
        if poly.material_index == target_material_index:
            poly.material_index = new_material_index
    
    print(f"已将 {len(selected_faces)} 个面分配给新的纹理材质")
    
    # 切换到材质预览模式
    for area in bpy.context.screen.areas:
        if area.type == 'VIEW_3D':
            area.spaces[0].shading.type = 'MATERIAL'
            break
    
    print("=== 脚本执行完成 ===")
    
    # 创建输出目录
    if output_path and os.path.dirname(output_path):
        os.makedirs(os.path.dirname(output_path), exist_ok=True) 
    
    # 导出结果
    if output_path:
        output_ext = os.path.splitext(output_path)[1].lower()
        
        try:
            if output_ext == '.obj':
                bpy.ops.wm.obj_export(filepath=output_path)
                print(f"已保存处理后的模型为OBJ格式: {output_path}")
                
            elif output_ext == '.fbx':
                bpy.ops.export_scene.fbx(filepath=output_path, use_selection=True)
                print(f"已保存处理后的模型为FBX格式: {output_path}")
                
            elif output_ext in ['.glb', '.gltf']:
                bpy.ops.export_scene.gltf(
                    filepath=output_path,
                    use_selection=True
                )
                print(f"已保存处理后的模型为{output_ext.upper()}格式: {output_path}")
                
            else:
                print(f"不支持的输出格式: {output_ext}，默认保存为.blend")
                output_path = os.path.splitext(output_path)[0] + '.blend'
                bpy.ops.wm.save_as_mainfile(filepath=output_path)
                
        except Exception as e:
            print(f"导出失败: {e}")
            output_path = os.path.splitext(output_path)[0] + '.blend'
            bpy.ops.wm.save_as_mainfile(filepath=output_path)
        
        print(f"最终保存路径: {output_path}")
    else:
        print("警告: 未指定输出路径，结果未保存")

def get_materials_info():
    """获取所有材质信息"""
    materials = []
    processed_materials = set()  # 避免重复
    
    for obj in bpy.data.objects:
        if obj.type == 'MESH':
            for i, slot in enumerate(obj.material_slots):
                if slot.material and slot.material.name not in processed_materials:
                    materials.append({
                        "name": slot.material.name,
                        "object": obj.name
                    })
                    processed_materials.add(slot.material.name)
    
    return materials

def export_materials_info(output_path):
    """导出材质信息到JSON文件"""
    materials = get_materials_info()
    
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(materials, f, indent=2, ensure_ascii=False)
        print(f"材质信息已导出到: {output_path}")
        print(f"找到 {len(materials)} 个材质:")
        for mat in materials:
            print(f"  - {mat['name']} (来自对象: {mat['object']})")
    except Exception as e:
        print(f"导出材质信息失败: {e}")

def texture_painting_interface(
    obj_path, 
    output_path, 
    texture_path,
    material_name=None
):
    """
    封装的纹理绘制接口函数
    :param obj_path: 输入模型文件路径
    :param output_path: 输出文件路径
    :param texture_path: 纹理图像路径
    :param material_name: 材质名称（用于部分区域选择）
    """
    setup_texture_painting(
        obj_path=obj_path,
        output_path=output_path,
        texture_path=texture_path,
        material_name=material_name
    )

if __name__ == "__main__":
    # 解析命令行参数
    try:
        script_args_start = sys.argv.index('--') + 1
        script_args = sys.argv[script_args_start:]
    except ValueError:
        script_args = []
    
    parser = argparse.ArgumentParser(description='Blender部分区域纹理绘制脚本 - 支持多种3D格式')
    parser.add_argument('--obj_path', default=r"C:\Users\86184\Desktop\AreaTest.obj", 
                       help='输入模型文件路径 (支持: .obj, .glb, .gltf)')
    parser.add_argument('--output_path', default=r"test_result_obj10.obj", 
                       help='输出文件路径 (支持: .obj, .glb, .gltf)')
    parser.add_argument('--texture_path', default=r"C:\Users\86184\Pictures\jujingyi.jpg", 
                       help='纹理图像路径')
    parser.add_argument('--material_name', default="ProjectionArea", 
                       help='材质名称（用于部分区域选择）')
    parser.add_argument('--get_materials', default="false", 
                       help='是否只获取材质信息')
    
    args = parser.parse_args(script_args)
    
    print("=== 部分区域纹理绘制脚本启动 ===")
    print(f"输入文件: {args.obj_path}")
    print(f"输出文件: {args.output_path}")

    # 检查是否只需要获取材质信息
    if args.get_materials.lower() == 'true':
        print("=== 材质信息获取模式 ===")
        
        # 清除当前场景中的所有对象
        bpy.ops.object.select_all(action='SELECT')
        bpy.ops.object.delete()
        
        # 导入3D模型
        try:
            import_3d_model(args.obj_path)
            print("模型导入成功，开始获取材质信息...")
            
            # 导出材质信息
            export_materials_info(args.output_path)
            
        except Exception as e:
            print(f"获取材质信息失败: {e}")
            # 输出空的材质列表，避免前端报错
            with open(args.output_path, 'w', encoding='utf-8') as f:
                json.dump([], f)
        
        print("=== 材质信息获取完成 ===")
    else:
        print("=== 纹理绘制模式 ===")
        print(f"纹理文件: {args.texture_path}")
        print(f"目标材质: {args.material_name}")
        print("=" * 50)
        
    # print(f"纹理文件: {args.texture_path}")
    # print(f"目标材质: {args.material_name}")
    # print("=" * 50)

    texture_painting_interface(
        obj_path=args.obj_path,
        output_path=args.output_path,
        texture_path=args.texture_path,
        material_name=args.material_name
    )

    # 代码运行 blender --background --python analyze_bound.py